import torch
import time


def gpu_stress_test(duration=5):
    # 检查GPU是否可用
    if not torch.cuda.is_available():
        print("没有可用的GPU，无法进行测试")
        return

    # 设置使用的GPU
    device = torch.device("cuda:0")
    print(f"使用GPU: {torch.cuda.get_device_name(device)}")

    # 创建大型随机向量
    vector_size = 2 ** 24  # 约1600万个元素
    num_vectors = 32  # 向量数量

    print(f"创建 {num_vectors} 个大小为 {vector_size} 的向量...")

    # 在GPU上创建随机向量
    vectors = [torch.randn(vector_size, device=device) for _ in range(num_vectors)]

    # 开始时间
    start_time = time.time()
    end_time = start_time + duration

    # 持续进行向量运算直到达到指定时长
    print(f"开始进行 {duration} 秒的向量计算...")
    operations = 0

    while time.time() < end_time:
        # 随机选择两个向量进行点积运算
        a = vectors[torch.randint(0, num_vectors, (1,)).item()]
        b = vectors[torch.randint(0, num_vectors, (1,)).item()]

        # 点积运算
        dot_product = torch.dot(a, b)

        # 向量加法
        c = a + b

        # 向量乘法（元素级）
        d = a * b

        # 向量缩放
        e = 0.5 * a + 0.3 * b

        operations += 1

    # 计算性能指标
    elapsed = time.time() - start_time
    print(f"测试完成！在 {elapsed:.2f} 秒内执行了 {operations} 组运算")
    print(f"GPU使用率应该已经短暂上升")


if __name__ == "__main__":
    # 运行5秒钟的GPU向量计算
    gpu_stress_test(duration=30)
